#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@File    : mngs_batch_plot_coverage.py
@Author  : Bing Liang
@Email   : believer19940901@gmail.com
@Date    : 2025/10/22 11:23
@Description : 系统对接 - 批量绘制物种覆盖率图
"""

from argparse import ArgumentParser
from pathlib import Path
from multiprocessing import Pool
import pandas as pd
import subprocess
import shutil


# ==================================================
# 全局路径配置
# ==================================================
PYTHON3 = Path("/home/bioinfo/software/miniconda3/envs/python-venv/bin/python3")
BIOTOOLS = Path("/data/mNGS/pipeline/script/autoReport/biotools/main.py")
PLOT = Path("/data/mNGS/pipeline/sys_api/mngs_scripts/plotDepthCoverage4flask.py")
REF_LENGTH = Path("/data/mNGS/pipeline/script/table/ref_length_4_coverage.xls")
WORK_DIR = Path("/data/mNGS/runmngs/tmp/batchPlotCoverage")
RESULT_DIR = Path("/data/mNGS/runmngs/result")


# ==================================================
# 辅助函数
# ==================================================
def _run_or_dead(cmd: str):
    """
    执行系统命令，如执行失败则终止程序。

    Args:
        cmd (str): 要执行的命令。
    """
    try:
        print(f"[INFO] CMD START: {cmd}")
        result = subprocess.run(
            cmd,
            shell=True,
            check=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True
        )
        print(f"[INFO] CMD FINISH: {cmd}")
        if result.returncode != 0:
            print(f"[ERROR] Command failed with code {result.returncode}")
            raise RuntimeError(f"Command failed: {cmd}")
    except subprocess.CalledProcessError as err:
        print(f"[ERROR] Command execution failed: {cmd}")
        print(f"[STDERR] {err.stderr.strip()}")
        raise


def _get_tenant_id(id_file: Path) -> int:
    """
    从文件中读取租户编号，若文件不存在或读取失败则返回 1。

    Args:
        id_file (Path): 租户编号文件路径。

    Returns:
        int: 租户编号。
    """
    try:
        if not id_file.exists() or not id_file.is_file():
            return 1
        with open(id_file, "r") as f:
            return int(f.read().strip())
    except Exception as err:
        print(f"[WARN] Failed to read tenant_id from {id_file}: {err}")
        return 1


def _get_sample_ai_by_jobid(jobid: str, tenant_id: int, outfile: Path):
    """
    根据芯片编号和租户编号获取检出结果。

    Args:
        jobid (str): 芯片编号。
        tenant_id (int): 租户编号。
        outfile (Path): 输出 TSV 文件路径。
    """
    cmd = (
        f"{PYTHON3} {BIOTOOLS} get_sample_ai_by_jobid "
        f"--jobid {jobid} "
        f"--tenant_id {tenant_id} "
        f"--out_tsv {outfile}"
    )
    _run_or_dead(cmd)


def _batch_plot_coverage(jobid: str, sample_ai: Path):
    """
    批量绘制样本覆盖率图。

    Args:
        jobid (str): 芯片编号。
        sample_ai (Path): 检出结果文件路径。
    """
    cmds = []
    jobid_out_dir = WORK_DIR / jobid
    if jobid_out_dir.exists():
        shutil.rmtree(jobid_out_dir)
    jobid_out_dir.mkdir(parents=True, exist_ok=True)

    df = pd.read_csv(sample_ai, sep="\t", encoding="utf-8")

    for _, row in df.iterrows():
        sample_id = row["sample_id"]

        # 仅处理高等级检出
        if row["species_focous"] == "高" and row["species_status"] == 1:
            species_name = row["species"].strip().split("|")[1].replace(" ", "_")
            lite_sam = RESULT_DIR / jobid / sample_id / f"{sample_id}.litesam4ref"
            species_out_file = jobid_out_dir / f"{sample_id}.{species_name}.png"

            cmd = (
                f"{PYTHON3} {PLOT} -i {sample_id} "
                f"-s {species_name} -a {REF_LENGTH} "
                f"-l {lite_sam} -o {species_out_file}"
            )
            cmds.append(cmd)

    # 并行执行绘图命令
    if cmds:
        with Pool(len(cmds)) as pool:
            for cmd in cmds:
                pool.apply_async(_run_or_dead, args=(cmd,))
            pool.close()
            pool.join()

    # 打包结果文件
    out_zip = WORK_DIR / f"{jobid}.zip"
    if out_zip.exists():
        out_zip.unlink()


    zip_cmd = f"cd {jobid_out_dir} && zip -r ../{jobid}.zip ./*.png"
    _run_or_dead(zip_cmd)

    print(f"[INFO] 绘图完成: {out_zip}")


# ==================================================
# 主流程
# ==================================================
def main(args):
    """
    主入口：获取检出结果并批量绘制覆盖率图。
    """
    jobid = args.job_id
    sample_ai = WORK_DIR / f"{jobid}.sample_ai.tsv"
    tenant_file = Path("/data/mNGS/runmngs/csv") / f"{jobid}.id"

    tenant_id = _get_tenant_id(tenant_file)
    _get_sample_ai_by_jobid(jobid, tenant_id, sample_ai)
    _batch_plot_coverage(jobid, sample_ai)


if __name__ == "__main__":
    parser = ArgumentParser(description="批量绘制 mNGS 样本覆盖率图")
    parser.add_argument("--job_id", type=str, required=True, help="芯片编号")
    main(parser.parse_args())
